package com.wisecoach.gatewayplus.transaction.lockedtx;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * {@code @author:} wisecoach
 * {@code @date:} 2023/9/26 上午10:55
 * {@code @version:} 1.0.0
 */
public class MemoryLockSource implements LockSource {

    private final Map<String, Map<String, RWLock>> source = new HashMap<>();
    private final Map<String, Lock> nsLocks = new ConcurrentHashMap<>();

    @Override
    public boolean lock(TxRWSet rwSet) {
        AtomicBoolean allLocked = new AtomicBoolean(true);
        // 可以保证每个nsSource线程安全
        List<Lock> locks = lockNs(rwSet.getReadSet().keySet());
        // Map<String, Set<String>> readSetToLock = getReadSetToLock(rwSet);
        // 先判断能否对写集上锁
        rwSet.getWriteSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = source.computeIfAbsent(ns, k -> new HashMap<>());
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                if (!rwLock.canLock()) {
                    allLocked.compareAndSet(true, false);
                }
            });
        });
        rwSet.getReadSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = source.computeIfAbsent(ns, k -> new HashMap<>());
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                if (!rwLock.canRLock()) {
                    allLocked.compareAndSet(true, false);
                }
            });
        });
        // 如果无法全部上锁
        if (!allLocked.get()) {
            locks.forEach(Lock::unlock);
            return false;
        }
        // 以下正式对所有要上锁的数据进行上锁
        rwSet.getWriteSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = source.computeIfAbsent(ns, k -> new HashMap<>());
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                boolean tryLock = rwLock.tryLock();
                if (!tryLock) {
                    throw new RuntimeException("出了些意外");
                }
            });
        });
        rwSet.getReadSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = source.computeIfAbsent(ns, k -> new HashMap<>());
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                boolean tryLock = rwLock.tryRLock();
                if (!tryLock) {
                    throw new RuntimeException("出了些意外");
                }
            });
        });
        // 释放nsSource的lock
        locks.forEach(Lock::unlock);
        return allLocked.get();
    }

    @Override
    public void unLock(TxRWSet rwSet) {
        // 可以保证每个nsSource线程安全
        List<Lock> locks = lockNs(rwSet.getReadSet().keySet());
        // Map<String, Set<String>> readSetToLock = getReadSetToLock(rwSet);
        rwSet.getWriteSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = source.computeIfAbsent(ns, k -> new HashMap<>());
            keys.forEach(key -> nsMap.get(key).unlock());
        });
        rwSet.getReadSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = source.computeIfAbsent(ns, k -> new HashMap<>());
            keys.forEach(key -> nsMap.get(key).unRLock());
        });
        // 释放nsSource的lock
        locks.forEach(Lock::unlock);
    }

    private List<Lock> lockNs(Set<String> nsSet) {
        // 如果namespace 只有一个，即不存在链码互相调用
        if (nsSet.size() == 1) {
            ArrayList<Lock> locks = new ArrayList<>();
            String ns = nsSet.iterator().next();
            nsLocks.compute(ns, (k, v) -> {
                if (v == null) {
                    ReentrantLock lock = new ReentrantLock();
                    lock.lock();
                    locks.add(lock);
                    return lock;
                } else {
                    v.lock();
                    locks.add(v);
                    return v;
                }
            });
            return locks;
        } else {
            synchronized (nsLocks) {
                ArrayList<Lock> locks = new ArrayList<>();
                nsSet.forEach(ns -> nsLocks.compute(ns, (k, v) -> {
                    if (v == null) {
                        ReentrantLock lock = new ReentrantLock();
                        lock.lock();
                        locks.add(lock);
                        return lock;
                    } else {
                        v.lock();
                        locks.add(v);
                        return v;
                    }
                }));
                return locks;
            }
        }
    }

    private Map<String, Set<String>> getReadSetToLock(TxRWSet rwSet) {
        HashMap<String, Set<String>> ret = new HashMap<>(rwSet.getReadSet());
        rwSet.getWriteSet().keySet().forEach(ret::remove);
        return ret;
    }

}
